import numpy as np
import soundfile as sf


def embed_watermark(audio_path, watermark, output_path):
    # 读取并归一化音频（关键修复点）
    audio, sample_rate = sf.read(audio_path)
    audio = np.clip(audio.mean(axis=1) if audio.ndim > 1 else audio, -1.0, 1.0)

    # 水印处理
    watermark_bin = ''.join(format(ord(c), '08b') for c in watermark)
    watermark_len = len(watermark_bin)
    if watermark_len > len(audio):
        raise ValueError("水印过长，音频容量不足")

    # 安全转为int16（关键修复点）
    audio_int = np.clip(audio  * (2**15), -32768, 32767).astype(np.int16)

    # LSB嵌入
    for i in range(watermark_len):
        if i >= len(audio_int):
            break
        if(audio_int[i] & 0xFE >= 32767):
            continue
        audio_int[i] = (audio_int[i] & 0xFE) | int(watermark_bin[i])

    sf.write(output_path, audio_int.astype(np.float32) / 32767, sample_rate)


# 测试（当前时间水印）
embed_watermark("original.wav", "2025-03-30-20:24", "output.wav")